package tapioca.util;

import java.util.ConcurrentModificationException;

import com.google.appengine.api.datastore.DatastoreFailureException;
import com.google.appengine.api.datastore.DatastoreService;
import com.google.appengine.api.datastore.DatastoreTimeoutException;
import com.google.appengine.api.datastore.Transaction;

public class RetryHelper {

  private final DatastoreService ds;
  private final WorkUnit workUnit;
  
  public static abstract class WorkUnit {
    public abstract void doWork(DatastoreService ds);
  }
  
  public static abstract class TransactionalWorkUnit extends WorkUnit {
    public void doWork(DatastoreService ds) {
      throw new UnsupportedOperationException(
          "Transactional work unit must be executed in a transaction");
    }
    
    public abstract void doWork(DatastoreService ds, Transaction txn);
  }
  
  public RetryHelper(DatastoreService ds, WorkUnit workUnit) {
    this.ds = ds;
    this.workUnit = workUnit;
  }
  
  public final void execute() {
    Transaction txn = null;
    Boolean retry = null;
    
    while (retry == null || retry == true) {
      // At the beginning of every loop, assume that we will succeed and that a retry is unneeded.      
      retry = false;
      
      try {
        if (workUnit instanceof TransactionalWorkUnit) {
          TransactionalWorkUnit txnWorkUnit = (TransactionalWorkUnit) workUnit;
          txn = ds.beginTransaction();
          txnWorkUnit.doWork(ds, txn);
          if (txn.isActive()) {
            txn.commit();
          }
        } else {
          workUnit.doWork(ds);
        }
        
      // These are the only "acceptable" retry exceptions currently.  Any other RuntimeException
      // such as DatastoreNeedIndexException are not valid for retries.
      // TODO (earmbrust): Check with maxr and ryanb to see if this is correct.
      } catch (ConcurrentModificationException e) {
        retry = true;
      } catch (DatastoreTimeoutException e) {
        retry = true;
      } catch (DatastoreFailureException e) {
        retry = true;
      } finally {
        // If the current transaction is active, this is a transaction that did not succeed so it
        // should be rolled back.
        if (txn != null && txn.isActive()) {
          txn.rollback();
        }
      }        
    }
  }
  
}
